"""DHF模型基本使用示例

演示如何使用DHF模型进行降雨径流模拟。
"""

import os
import sys
import time

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
from loguru import logger

from service.dhf_model import DHFModel

import warnings
import matplotlib

# 抑制所有matplotlib警告
warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib')
matplotlib.set_loglevel('ERROR')


def load_rainfall_from_your_json(json_file_path: str, current_time: datetime):
    """
    从您的JSON文件格式加载降雨数据，以当前时间自动划分预热期和预报期

    Args:
        json_file_path: JSON文件路径
        current_time: 当前时间（作为预热期和预报期的分界点）

    Returns:
        包含预热期和预报降雨数据的字典
    """
    import json
    from datetime import datetime, timedelta

    logger.info(f"从JSON文件加载降雨数据: {json_file_path}")
    logger.info(f"当前时间分界点: {current_time}")

    try:
        with open(json_file_path, 'r', encoding='utf-8') as f:
            rainfall_data = json.load(f)

        # 解析所有降雨数据
        all_records = []

        for record in rainfall_data:
            try:
                # 解析时间
                dt_str = record.get('dt', '')
                record_time = datetime.strptime(dt_str, '%Y-%m-%d %H:%M:%S')

                # 获取降雨量，优先使用 rainReal，如果不存在或为0则使用 rain
                # rain_real = record.get('rainReal', 0)
                rain = record.get('rain', 0)

                if record_time < current_time:
                    rainfall_value = rain
                else:
                    rainfall_value = rain

                all_records.append({
                    'time': record_time,
                    'rainfall': max(0, rainfall_value),  # 确保非负
                    'original_record': record
                })

            except ValueError as e:
                logger.warning(f"跳过无效时间记录: {dt_str}, 错误: {e}")
                continue

        if not all_records:
            raise ValueError("未找到有效的降雨数据")

        # 按时间排序
        all_records.sort(key=lambda x: x['time'])

        # 以当前时间为分界点，划分预热期和预报期
        warmup_records = []
        forecast_records = []

        for record in all_records:
            if record['time'] < current_time:
                warmup_records.append(record)
            else:
                forecast_records.append(record)

        # 提取数据
        warmup_rainfall = [r['rainfall'] for r in warmup_records]
        warmup_times = [r['time'] for r in warmup_records]
        forecast_rainfall = [r['rainfall'] for r in forecast_records]
        forecast_times = [r['time'] for r in forecast_records]

        # 合并所有数据
        combined_rainfall = warmup_rainfall + forecast_rainfall
        combined_times = warmup_times + forecast_times

        # 计算模拟开始时间
        if warmup_times:
            simulation_start_time = warmup_times[0]
        elif forecast_times:
            simulation_start_time = forecast_times[0]
        else:
            simulation_start_time = current_time

        # 计算统计信息
        warmup_hours = len(warmup_rainfall)
        forecast_hours = len(forecast_rainfall)
        total_hours = len(combined_rainfall)
        current_time_index = warmup_hours  # 当前时间在合并数据中的索引位置

        result = {
            'combined_rainfall': combined_rainfall,
            'warmup_rainfall': warmup_rainfall,
            'forecast_rainfall': forecast_rainfall,
            'simulation_start_time': simulation_start_time,
            'current_time': current_time,
            'current_time_index': current_time_index,
            'warmup_hours': warmup_hours,
            'forecast_hours': forecast_hours,
            'total_hours': total_hours,
            'time_series': combined_times,
            'warmup_times': warmup_times,
            'forecast_times': forecast_times,
            'original_data_count': len(rainfall_data),
            'valid_data_count': len(all_records)
        }

        logger.info(f"降雨数据处理完成:")
        logger.info(f"  原始数据记录: {len(rainfall_data)} 条")
        logger.info(f"  有效数据记录: {len(all_records)} 条")
        logger.info(f"  预热期: {warmup_hours}小时 (当前时间之前)")
        logger.info(
            f"    - 时间范围: {warmup_times[0] if warmup_times else '无'} 至 {warmup_times[-1] if warmup_times else '无'}")
        logger.info(f"    - 总降雨量: {sum(warmup_rainfall):.1f}mm")
        logger.info(f"  预报期: {forecast_hours}小时 (当前时间之后)")
        logger.info(
            f"    - 时间范围: {forecast_times[0] if forecast_times else '无'} 至 {forecast_times[-1] if forecast_times else '无'}")
        logger.info(f"    - 总降雨量: {sum(forecast_rainfall):.1f}mm")
        logger.info(f"  模拟开始时间: {simulation_start_time}")
        logger.info(f"  当前时间分界点: {current_time}")

        return result

    except FileNotFoundError:
        logger.error(f"找不到降雨数据文件: {json_file_path}")
        raise
    except json.JSONDecodeError as e:
        logger.error(f"JSON文件格式错误: {e}")
        raise
    except Exception as e:
        logger.error(f"加载降雨数据时出错: {e}")
        raise


def create_sample_rainfall_data():
    """创建示例降雨数据"""
    # 模拟一次洪水过程的降雨数据 (72小时)
    hours = 72
    rainfall = np.zeros(hours)

    # 设计一个典型的暴雨过程
    # 前期小雨
    rainfall[0:12] = np.random.uniform(0, 2, 12)

    # 主雨期 (12-36小时)
    peak_hours = np.arange(12, 36)
    for i, hour in enumerate(peak_hours):
        if 18 <= hour <= 30:  # 暴雨中心
            rainfall[hour] = 15 + 10 * np.sin(np.pi * (hour - 18) / 12) + np.random.uniform(-2, 2)
        else:
            rainfall[hour] = 5 + np.random.uniform(-1, 2)

    # 后期退水
    rainfall[36:48] = np.random.uniform(0, 3, 12)
    rainfall[48:] = np.random.uniform(0, 1, hours - 48)

    # 确保非负值
    rainfall = np.maximum(rainfall, 0)

    return rainfall.tolist()


def create_sample_observed_flow():
    """创建示例观测流量数据 (用于模型验证)"""
    # 这里创建一个理想化的流量过程，实际应用中应使用真实观测数据
    hours = 120  # 比降雨数据长一些，包含完整的汇流过程

    # 基流
    base_flow = 50.0
    flow = np.full(hours, base_flow)

    # 洪峰过程 (简化模拟)
    peak_start = 20
    peak_end = 60
    peak_flow = 800.0

    for i in range(peak_start, peak_end):
        if i < 35:
            # 涨水段
            flow[i] = base_flow + (peak_flow - base_flow) * ((i - peak_start) / 15) ** 1.5
        else:
            # 退水段
            flow[i] = peak_flow * np.exp(-0.1 * (i - 35))

    # 添加一些随机噪声
    noise = np.random.normal(0, flow * 0.05)
    flow = np.maximum(flow + noise, base_flow * 0.8)

    return flow.tolist()


def main():
    """主函数"""
    logger.info("DHF模型基本使用示例")
    logger.info("=" * 50)

    # 1. 创建模型实例
    logger.info("\n1. 创建DHF模型实例...")
    model = DHFModel(
        initial_flow=50.0,  # 起始流量 50 m³/s
        rainfall_center='中游',  # 降雨中心位置
        watershed_area=1200.0,  # 流域面积 1200 km²
        config_file='./output/calibrated_parameters.json'  # 模型配置文件config_file
    )
    logger.info(f"模型信息: {model}")

    # 2. 准备输入数据
    logger.info("\n2. 准备输入数据...")
    rainfall_data = create_sample_rainfall_data()
    start_time = datetime(2023, 7, 15, 8, 0)  # 2023年7月15日8时开始
    time_step = 1
    # rainfall_data_info = load_rainfall_from_your_json("./8.json", datetime(2025, 8, 6, 14, 0))
    # rainfall_data = rainfall_data_info['combined_rainfall']
    # import pandas as pd
    # start_time = pd.to_datetime(rainfall_data_info['warmup_times'][0])
    # time_step = 3

    logger.info(f"降雨数据长度: {len(rainfall_data)} 小时")
    logger.info(f"总降雨量: {sum(rainfall_data):.1f} mm")
    logger.info(f"最大小时降雨量: {max(rainfall_data):.1f} mm")

    # 3. 运行模型模拟
    logger.info("\n3. 运行模型模拟...")
    results = model.run_simulation(
        rainfall_data=rainfall_data,
        time_step=time_step,  # 1小时时间步长
        start_time=start_time
    )
    logger.warning(results)
    logger.info(f"模拟结果时间范围: {results['time_series'][0]} 至 {results['time_series'][-1]}")

    # 4. 显示模拟结果
    logger.info("\n4. 模拟结果:")
    peak_info = results['peak_info']
    logger.info(f"洪峰流量: {peak_info['peak_flow']:.2f} m³/s")
    logger.info(f"峰现时间: {peak_info['peak_time']:.1f} 小时")
    logger.info(f"径流总量: {results['total_volume']:.2f} 万m³")

    # 5. 模型性能评估 (使用模拟的观测数据)
    logger.info("\n5. 模型性能评估...")
    observed_flow = create_sample_observed_flow()
    performance = model.evaluate_performance(observed_flow)

    logger.info("性能指标:")
    for metric, value in performance.items():
        logger.info(f"  {metric}: {value:.4f}")

    # 6. 生成图表
    import os
    logger.info("\n6. 生成结果图表...")
    try:
        # 创建输出目录（在生成图表之前）
        os.makedirs("./output", exist_ok=True)

        plots = model.generate_plots(save_dir="./output")
        logger.success("图表已生成并保存到 ./output 目录")

        # 确保字体配置对交互式显示生效
        from util.font_config import setup_chinese_font
        setup_chinese_font()

        # 显示图表
        plt.show()

    except Exception as e:
        logger.error(f"图表生成出现问题: {e}")
        logger.warning("请确保安装了matplotlib和相关依赖")

    # 7. 导出结果
    logger.info("\n7. 导出结果...")
    try:
        # 确保输出目录存在（防御性编程）
        os.makedirs("./output", exist_ok=True)

        # 导出降雨数据到Excel
        import pandas as pd

        # 创建时间序列
        time_index = pd.date_range(start=start_time, periods=len(rainfall_data), freq='H')

        # 创建降雨数据DataFrame
        rainfall_df = pd.DataFrame({
            '时间': time_index,
            '降雨量(mm)': rainfall_data
        })

        # 导出降雨数据
        rainfall_df.to_excel("./output/rainfall_input_data.xlsx", index=False, sheet_name='降雨数据')
        logger.success("降雨输入数据已导出到 ./output/rainfall_input_data.xlsx")

        # 导出Excel文件
        model.export_results("./output/dhf_model_results.xlsx")
        logger.success("结果已导出到 ./output/dhf_model_results.xlsx")

        # 生成报告
        report = model.generate_report()
        with open("./output/model_report.txt", "w", encoding="utf-8") as f:
            f.write(report)
        logger.success("模型报告已保存到 ./output/model_report.txt")

    except Exception as e:
        logger.error(f"结果导出出现问题: {e}")

    # 8. 参数率定示例 (可选)
    logger.info("\n8. 参数率定示例...")

    # 短暂延迟确保日志输出完成
    time.sleep(0.1)

    response = input("是否进行参数率定? (y/n): ")
    if response.lower() == 'y':
        logger.info("开始参数率定 (这可能需要几分钟时间)...")

        # 重置模型状态
        model.reset_model()

        # 进行参数率定
        calibration_results = model.calibrate_parameters(
            rainfall_data=rainfall_data,
            observed_flow=observed_flow[:len(results['flow_series'])],
            max_iterations=50
        )

        logger.info("\n率定结果:")
        if calibration_results['best_parameters']:
            logger.info("最佳参数:")
            for param, value in calibration_results['best_parameters'].items():
                logger.info(f"  {param}: {value:.4f}")
            logger.info(f"最佳NSE: {calibration_results['best_nse']:.4f}")

            # 使用优化后的参数重新运行模拟
            logger.info("\n使用优化参数重新运行模拟...")
            optimized_results = model.run_simulation(
                rainfall_data=rainfall_data,
                time_step=1.0,
                start_time=start_time
            )

            # 显示优化后的模拟结果
            logger.info("\n优化后模拟结果:")
            optimized_peak_info = optimized_results['peak_info']
            logger.info(f"洪峰流量: {optimized_peak_info['peak_flow']:.2f} m³/s")
            logger.info(f"峰现时间: {optimized_peak_info['peak_time']:.1f} 小时")
            logger.info(f"径流总量: {optimized_results['total_volume']:.2f} 万m³")

            # 对比优化前后的效果
            logger.info("\n优化前后对比:")
            original_peak = results['peak_info']['peak_flow']
            optimized_peak = optimized_peak_info['peak_flow']
            peak_improvement = ((optimized_peak - original_peak) / original_peak) * 100
            logger.info(f"洪峰流量变化: {original_peak:.2f} → {optimized_peak:.2f} m³/s ({peak_improvement:+.1f}%)")

            original_volume = results['total_volume']
            optimized_volume = optimized_results['total_volume']
            volume_improvement = ((optimized_volume - original_volume) / original_volume) * 100
            logger.info(
                f"径流总量变化: {original_volume:.2f} → {optimized_volume:.2f} 万m³ ({volume_improvement:+.1f}%)")

            # 重新评估性能并生成对比图表
            optimized_performance = model.evaluate_performance(observed_flow)
            logger.info(f"\n性能指标对比:")
            logger.info(f"NSE: {calibration_results['best_nse']:.4f}")
            logger.info(f"R²: {optimized_performance['R2']:.4f}")
            logger.info(f"RMSE: {optimized_performance['RMSE']:.2f}")

            # 生成优化后的图表
            logger.info("\n生成优化后的图表...")
            import os  # 如果文件顶部没有导入的话
            os.makedirs("./output/optimized", exist_ok=True)
            optimized_plots = model.generate_plots(save_dir="./output/optimized")
            logger.success("优化后图表已保存到 ./output/optimized/")

        else:
            logger.warning("参数率定未能改善模型性能")

    logger.success("\n示例运行完成！")
    logger.info("\n模型使用说明:")
    logger.info("1. 根据实际流域调整模型参数 (流域面积、降雨中心等)")
    logger.info("2. 使用真实的降雨和蒸发数据")
    logger.info("3. 利用历史洪水数据进行参数率定")
    logger.info("4. 验证模型在不同洪水事件上的表现")


if __name__ == "__main__":
    main()
